"""
LangGraph实现的初步审查报告生成流程
包含三个节点：
1. 节点1：解析文件（OCR识别）
2. 节点2：调用大模型处理
3. 节点3：结束
"""
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict
from typing import Annotated, Optional
from pathlib import Path
import asyncio
import sys
import os
import yaml

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.insert(0, project_root)

from call_llm import _run_all_configs_once

# 加载配置文件
CONFIG = {}
try:
    config_path = os.path.join(project_root, "config.yaml")
    if os.path.exists(config_path):
        with open(config_path, "r", encoding="utf-8") as f:
            CONFIG = yaml.safe_load(f) or {}
            if not isinstance(CONFIG, dict):
                CONFIG = {}
except Exception as e:
    print(f"加载 config.yaml 失败：{e}")

# 从配置文件读取prompt配置
LLM_PROMPTS_CONFIG = CONFIG.get("LLM_PROMPTS", {})
LLM1_PROMPTS = LLM_PROMPTS_CONFIG.get("llm_summarize", {})
LLM2_PROMPTS = LLM_PROMPTS_CONFIG.get("llm2", {})


def load_prompt_from_file(file_path: str) -> Optional[str]:
    """
    从外部文件加载提示词
    
    Args:
        file_path: 文件路径（支持相对路径和绝对路径）
    
    Returns:
        str: 文件内容，如果文件不存在则返回None
    """
    try:
        # 如果是相对路径，相对于项目根目录
        if not os.path.isabs(file_path):
            file_path = os.path.join(project_root, file_path)
        
        if os.path.exists(file_path):
            with open(file_path, "r", encoding="utf-8") as f:
                return f.read().strip()
        return None
    except Exception as e:
        print(f"加载提示词文件失败 {file_path}: {e}")
        return None


def get_prompt_from_config(config: dict, prompt_key: str) -> Optional[str]:
    """
    从配置中获取提示词，支持直接字符串或文件路径
    
    Args:
        config: 配置字典
        prompt_key: 提示词的键名（如 'system_prompt' 或 'user_prompt'）
    
    Returns:
        str: 提示词内容，如果未找到则返回None
    """
    prompt_value = config.get(prompt_key)
    if not prompt_value:
        return None
    
    # 如果以 @file: 开头，表示从文件读取
    if isinstance(prompt_value, str) and prompt_value.startswith("@file:"):
        file_path = prompt_value[6:].strip()  # 去掉 "@file:" 前缀
        return load_prompt_from_file(file_path)
    
    # 否则直接返回字符串
    return prompt_value


class ReviewState(TypedDict):
    """状态定义"""
    file_path: Optional[str]  # 文件路径
    parsed_text: Optional[str]  # 解析后的文本内容
    file_type: Optional[str]  # 文件类型
    llm_result: Optional[str]  # 大模型1处理结果
    llm_result2: Optional[str]  # 大模型2处理结果
    error: Optional[str]  # 错误信息
    llm_name: Optional[str]  # 大模型1名称，默认 "qwen3-14b"
    llm_name2: Optional[str]  # 大模型2名称，默认 "qwen3-14b"
    system_prompt: Optional[str]  # 大模型1系统提示词
    system_prompt2: Optional[str]  # 大模型2系统提示词
    user_prompt: Optional[str]  # 大模型1用户提示词（如果提供，会替换自动生成的）
    user_prompt2: Optional[str]  # 大模型2用户提示词（如果提供，会替换自动生成的）
    concurrent: Optional[int]  # 并发数，默认 1


def parse_file_node(state: ReviewState) -> ReviewState:
    """
    节点1：解析上传的文件
    支持PDF和图片的OCR识别
    """
    try:
        file_path = state.get("file_path")
        if not file_path:
            return {"error": "文件路径为空"}
        
        file_path_obj = Path(file_path)
        if not file_path_obj.exists():
            return {"error": f"文件不存在: {file_path}"}
        
        # 获取临时目录
        tmp_dir = file_path_obj.parent
        
        # 解析文件
        parse_result = parse_file(file_path_obj, tmp_dir)
        
        return {
            "parsed_text": parse_result.get("full_text", ""),
            "file_type": parse_result.get("file_type", "unknown"),
            "error": None
        }
    except Exception as e:
        return {"error": f"文件解析失败: {str(e)}"}


def llm_summarize(state: ReviewState) -> ReviewState:
    """
    节点2：调用大模型将所有总结好的信访文件 最终汇总成一份提炼好的信访报告
    使用call_llm的_run_all_configs_once方法
    参数从state中获取，如果未提供则使用默认值
    """
    try:
        parsed_text = state.get("parsed_text")
        if not parsed_text:
            return {"error": "解析后的文本为空，无法调用大模型"}
        
        # 从state中获取参数，如果未提供则从配置文件读取，最后使用默认值
        llm_name = state.get("llm_name") or "qwen3-14b"
        system_prompt = state.get("system_prompt") or get_prompt_from_config(LLM1_PROMPTS, "system_prompt") or "你是一个专业的文档分析助手，请对上传的文档内容进行分析和总结。"
        user_prompt_template = state.get("user_prompt") or get_prompt_from_config(LLM1_PROMPTS, "user_prompt")
        concurrent = state.get("concurrent") or 1
        
        # 如果提供了自定义user_prompt，使用它；否则自动生成
        if user_prompt_template:
            # 如果user_prompt中包含{parsed_text}占位符，则替换；否则直接使用
            if "{parsed_text}" in user_prompt_template:
                try:
                    user_prompt = user_prompt_template.replace("{parsed_text}", parsed_text)
                except Exception:
                    # 如果替换失败，直接使用模板
                    user_prompt = user_prompt_template
            else:
                user_prompt = user_prompt_template
        else:
            # 默认的user_prompt
            user_prompt = f"请分析以下文档内容，并总结操作履历和线索：\n\n{parsed_text}"
        
        # 调用大模型（异步调用）
        try:
            # 尝试获取当前事件循环
            loop = asyncio.get_event_loop()
        except RuntimeError:
            # 如果没有事件循环，创建新的
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            should_close = True
        else:
            should_close = False
        
        try:
            results = loop.run_until_complete(
                _run_all_configs_once(
                    name=llm_name,
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    concurrent=concurrent
                )
            )
        finally:
            if should_close:
                loop.close()
        
        # 处理结果
        if results and len(results) > 0:
            llm_result = results[0] if isinstance(results[0], str) else str(results[0])
        else:
            llm_result = "大模型调用失败，未返回结果"
        
        return {
            "llm_result": llm_result,
            "error": None
        }
    except Exception as e:
        return {"error": f"大模型调用失败: {str(e)}"}


def call_llm_node2(state: ReviewState) -> ReviewState:
    """
    节点2：调用大模型2处理解析后的文本
    使用call_llm的_run_all_configs_once方法
    参数从state中获取，如果未提供则从配置文件读取，最后使用默认值
    """
    try:
        parsed_text = state.get("parsed_text")
        if not parsed_text:
            return {"error": "解析后的文本为空，无法调用大模型"}

        # 从state中获取参数，如果未提供则从配置文件读取，最后使用默认值
        llm_name = state.get("llm_name2") or "qwen3-14b"
        system_prompt = state.get("system_prompt2") or get_prompt_from_config(LLM2_PROMPTS, "system_prompt") or "你是一个专业的文档分析助手，请对上传的文档内容进行分析和总结。"
        user_prompt_template = state.get("user_prompt2") or get_prompt_from_config(LLM2_PROMPTS, "user_prompt")
        concurrent = state.get("concurrent") or 1

        # 如果提供了自定义user_prompt，使用它；否则自动生成
        if user_prompt_template:
            # 如果user_prompt中包含{parsed_text}占位符，则替换；否则直接使用
            if "{parsed_text}" in user_prompt_template:
                try:
                    user_prompt = user_prompt_template.replace("{parsed_text}", parsed_text)
                except Exception:
                    # 如果替换失败，直接使用模板
                    user_prompt = user_prompt_template
            else:
                user_prompt = user_prompt_template
        else:
            # 默认的user_prompt
            user_prompt = f"请分析以下文档内容，并总结操作履历和线索：\n\n{parsed_text}"

        # 调用大模型（异步调用）
        try:
            # 尝试获取当前事件循环
            loop = asyncio.get_event_loop()
        except RuntimeError:
            # 如果没有事件循环，创建新的
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            should_close = True
        else:
            should_close = False

        try:
            results = loop.run_until_complete(
                _run_all_configs_once(
                    name=llm_name,
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    concurrent=concurrent
                )
            )
        finally:
            if should_close:
                loop.close()

        # 处理结果
        if results and len(results) > 0:
            llm_result = results[0] if isinstance(results[0], str) else str(results[0])
        else:
            llm_result = "大模型调用失败，未返回结果"

        return {
            "llm_result2": llm_result,
            "error": None
        }
    except Exception as e:
        return {"error": f"大模型调用失败: {str(e)}"}

def end_node(state: ReviewState) -> ReviewState:
    """
    节点3：结束节点
    返回最终结果
    """
    return state


# 构建LangGraph
graph_builder = StateGraph(ReviewState)
graph_builder.add_node("parse_file", parse_file_node)
graph_builder.add_node("call_llm", call_llm_node)
graph_builder.add_node("call_llm2", call_llm_node2)
graph_builder.add_node("end", end_node)

# 定义边
graph_builder.add_edge(START, "parse_file")
graph_builder.add_edge("parse_file", "call_llm")
graph_builder.add_edge("parse_file", "call_llm2")
graph_builder.add_edge("call_llm", "end")
graph_builder.add_edge("call_llm2", "end")
graph_builder.add_edge("end", END)

# 编译图
graph = graph_builder.compile()


def visualize_graph(format: str = "mermaid") -> str:
    """
    可视化LangGraph
    
    Args:
        format: 可视化格式，可选 "mermaid", "ascii", "png", "pdf"
    
    Returns:
        str: 可视化结果（Mermaid代码、ASCII文本或文件路径）
    """
    try:
        graph_draw = graph.get_graph()
        
        if format == "mermaid":
            # 生成Mermaid格式的图（可以在支持Mermaid的Markdown编辑器中查看）
            mermaid_code = graph_draw.draw_mermaid()
            return mermaid_code
        elif format == "ascii":
            # 生成ASCII格式的图（在终端中查看）
            ascii_art = graph_draw.draw_ascii()
            return ascii_art
        elif format in ["png", "pdf"]:
            # 生成图片格式（需要安装graphviz）
            try:
                output_path = graph_draw.draw_mermaid_png()
                if format == "pdf":
                    # 如果需要PDF，可以转换或使用其他方法
                    return output_path.replace(".png", ".pdf") if output_path else None
                return output_path
            except Exception as e:
                return f"生成图片失败: {str(e)}，请确保已安装graphviz"
        else:
            return f"不支持的格式: {format}，支持: mermaid, ascii, png, pdf"
    except Exception as e:
        return f"可视化失败: {str(e)}"


def process_file(
    file_path: str,
    llm_name: Optional[str] = None,
    llm_name2: Optional[str] = None,
    system_prompt: Optional[str] = None,
    system_prompt2: Optional[str] = None,
    user_prompt: Optional[str] = None,
    user_prompt2: Optional[str] = None,
    concurrent: Optional[int] = None
) -> dict:
    """
    处理文件的入口函数
    
    Args:
        file_path: 文件路径
        llm_name: 大模型1名称，默认 "qwen3-14b"
        llm_name2: 大模型2名称，默认 "qwen3-14b"
        system_prompt: 大模型1系统提示词，默认从配置文件读取
        system_prompt2: 大模型2系统提示词，默认从配置文件读取
        user_prompt: 大模型1用户提示词，如果提供，可以使用{parsed_text}占位符，默认从配置文件读取
        user_prompt2: 大模型2用户提示词，如果提供，可以使用{parsed_text}占位符，默认从配置文件读取
        concurrent: 并发数，默认 1
    
    Returns:
        dict: 处理结果，包含llm_result和llm_result2
    """
    initial_state = {
        "file_path": file_path,
        "parsed_text": None,
        "file_type": None,
        "llm_result": None,
        "llm_result2": None,
        "error": None,
        "llm_name": llm_name,
        "llm_name2": llm_name2,
        "system_prompt": system_prompt,
        "system_prompt2": system_prompt2,
        "user_prompt": user_prompt,
        "user_prompt2": user_prompt2,
        "concurrent": concurrent
    }
    
    # 运行图并收集所有状态更新
    final_state = initial_state.copy()
    for event in graph.stream(initial_state):
        # 更新状态（每个节点都会更新状态）
        for node_name, node_output in event.items():
            if isinstance(node_output, dict):
                final_state.update(node_output)
    
    return final_state


# 测试可视化（如果直接运行此文件）
if __name__ == "__main__":
    print("=== ASCII 可视化 ===")
    print(visualize_graph("ascii"))
    print("\n=== Mermaid 代码 ===")
    print(visualize_graph("mermaid"))